package socket

import (
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"github.com/chicken-team-outside/chicken_transmission/api_messages"
	"github.com/chicken-team-outside/chicken_transmission/utils"
	"github.com/golang/protobuf/proto"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

var endianOrder = binary.LittleEndian

var byte4Pool = sync.Pool{New: func() any {
	return make([]byte, 4)
}}

func get4ByteBuf() []byte {
	return byte4Pool.Get().([]byte)
}

func readHeader(r io.Reader) (*api_messages.MessageHeader, error) {
	buf := get4ByteBuf()
	defer byte4Pool.Put(buf)
	_, err := io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}
	size := int32(endianOrder.Uint32(buf))
	buf = make([]byte, size)
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return nil, err
	}
	msg := new(api_messages.MessageHeader)
	err = proto.Unmarshal(buf, msg)
	if err != nil {
		return nil, err
	}
	return msg, nil
}

type msgPack struct {
	ctx     context.Context
	cancel  context.CancelCauseFunc
	msgType int
	data    io.Reader
}

var errSendSuccess = errors.New("send message success")

type OnMessageFunc func(ctx context.Context, header *api_messages.MessageHeader, body io.Reader)

type WebSocketConnection struct {
	conn         *websocket.Conn
	ctx          context.Context
	cancel       context.CancelCauseFunc
	lastPing     atomic.Int64
	lastPong     atomic.Int64
	once         sync.Once
	onClose      func()
	onMessage    OnMessageFunc
	maxSize      int
	msgWriteChan chan *msgPack
}

func (c *WebSocketConnection) GetCtx() context.Context {
	return c.ctx
}

func (c *WebSocketConnection) KeepAlive(checkTime time.Duration) {
	timeout := int64(checkTime.Seconds()) * 3
	tick := time.NewTimer(checkTime)
	defer tick.Stop()
	for {
		select {
		case <-c.ctx.Done():
			return
		case <-tick.C:
			if c.lastPong.Load()-c.lastPing.Load() > timeout {
				c.Close()
				return
			}
			c.lastPong.Store(time.Now().UnixMilli())
			err := c.ping()
			if err != nil {
				zap.L().Warn("failed to send ping message", zap.Error(err))
			} else {
				zap.L().Debug("sent ping message successfully")
			}
		}
	}
}

func (c *WebSocketConnection) ping() error {
	return c.SendRawMsg(context.Background(), websocket.PingMessage, nil)
}

func (c *WebSocketConnection) closeMsgWriteChan() {
	defer func() {
		if err := recover(); err != nil {
			zap.L().Debug("close message write channel occur panic", zap.Any("panic", err))
		}
	}()
	if c.msgWriteChan != nil {
		close(c.msgWriteChan)
	}
}

func (c *WebSocketConnection) Close() {
	c.once.Do(func() {
		c.cancel(ErrConnectionClosed)
		if err := c.conn.Close(); err != nil {
			zap.L().Debug("close connection failed", zap.Error(err))
		}
		c.closeMsgWriteChan()
		if c.onClose != nil {
			c.onClose()
		}
	})
}

func (c *WebSocketConnection) handlerWriteMessage() {
	defer c.Close()
	buf := make([]byte, 32*1024)
	for {
		select {
		case <-c.ctx.Done():
			return
		case msg := <-c.msgWriteChan:
			if msg.ctx.Err() == nil {
				var err error
				if msg.data == nil {
					err = c.conn.WriteMessage(msg.msgType, nil)
				} else {
					var writer io.WriteCloser
					writer, err = c.conn.NextWriter(msg.msgType)
					if err == nil {
						_, err = utils.CopyBuffer(msg.ctx, writer, msg.data, buf)
						if err01 := writer.Close(); err01 != nil {
							zap.L().Warn("failed to close message writer", zap.Error(err))
						}
					}
				}
				if err != nil {
					msg.cancel(err)
				} else {
					msg.cancel(errSendSuccess)
				}
			}
		}
	}
}

func (c *WebSocketConnection) SendRawMsgWithCtx(ctx, connCtx context.Context, msgType int, data io.Reader) (err error) {
	defer func() {
		e := recover()
		if e != nil {
			err = ErrConnectionClosed
		}
	}()
	msg := &msgPack{
		msgType: msgType,
		data:    data,
	}
	msg.ctx, msg.cancel = context.WithCancelCause(ctx)
	select {
	case <-connCtx.Done():
		return utils.GetCauseFromContext(connCtx)
	case c.msgWriteChan <- msg:
		select {
		case <-msg.ctx.Done():
			ctx = msg.ctx
		case <-connCtx.Done():
			return utils.GetCauseFromContext(connCtx)
		}
	case <-ctx.Done():
	}
	err = utils.GetCauseFromContext(ctx)
	if err != nil {
		if errors.Is(err, errSendSuccess) {
			return nil
		}
		return err
	}
	return err
}

func (c *WebSocketConnection) SendRawMsg(ctx context.Context, msgType int, data io.Reader) (err error) {
	return c.SendRawMsgWithCtx(ctx, c.ctx, msgType, data)
}

func (c *WebSocketConnection) SendMessageWithConnCtx(ctx context.Context, connCtx context.Context, header *api_messages.MessageHeader, body io.Reader) error {
	zap.L().Debug("send message", zap.Any("header", header))
	hData, err := proto.Marshal(header)
	if err != nil {
		return err
	}
	headerLen := make([]byte, 4)
	endianOrder.PutUint32(headerLen, uint32(len(hData)))
	if body == nil {
		return c.SendRawMsgWithCtx(ctx, connCtx, websocket.BinaryMessage, io.MultiReader(
			bytes.NewReader(headerLen),
			bytes.NewReader(hData),
		))
	} else {
		return c.SendRawMsgWithCtx(ctx, connCtx, websocket.BinaryMessage, io.MultiReader(
			bytes.NewReader(headerLen),
			bytes.NewReader(hData),
			body,
		))
	}
}

func (c *WebSocketConnection) SendMessage(ctx context.Context, header *api_messages.MessageHeader, body io.Reader) error {
	return c.SendMessageWithConnCtx(ctx, c.ctx, header, body)
}

func (c *WebSocketConnection) SendProtoMessageWithConnCtx(ctx context.Context, connCtx context.Context, header *api_messages.MessageHeader, body proto.Message) error {
	if body == nil {
		return c.SendMessageWithConnCtx(ctx, connCtx, header, nil)
	}
	bodyBytes, err := proto.Marshal(body)
	if err != nil {
		return err
	}
	header.BodySize = int32(len(bodyBytes))
	return c.SendMessageWithConnCtx(ctx, connCtx, header, bytes.NewReader(bodyBytes))
}

func (c *WebSocketConnection) SendProtoMessage(ctx context.Context, header *api_messages.MessageHeader, body proto.Message) error {
	return c.SendProtoMessageWithConnCtx(ctx, c.ctx, header, body)
}

func (c *WebSocketConnection) HandleMessage() error {
	defer c.Close()
	for c.ctx.Err() == nil {
		messageType, r, err := c.conn.NextReader()
		if err != nil {
			return err
		}
		if messageType == websocket.BinaryMessage {
			r = io.LimitReader(r, int64(c.maxSize))
			var header *api_messages.MessageHeader
			header, err = readHeader(r)
			if err != nil {
				zap.L().Warn("failed to read header from server", zap.Error(err))
			}
			if c.onMessage != nil {
				func() {
					defer func() {
						if e := recover(); e != nil {
							zap.L().Error("handle message occur panic", zap.Any("panic", e))
						}
					}()
					c.onMessage(c.ctx, header, r)
				}()
			}
		}
	}
	return ErrConnectionClosed
}

func NewWebSocket(conn *websocket.Conn, ctx context.Context, onClose func(), onMessage OnMessageFunc, maxMsgSize int) *WebSocketConnection {
	c := &WebSocketConnection{
		conn:         conn,
		onClose:      onClose,
		onMessage:    onMessage,
		maxSize:      maxMsgSize,
		msgWriteChan: make(chan *msgPack, 10),
	}
	c.ctx, c.cancel = context.WithCancelCause(ctx)
	conn.SetPingHandler(func(appData string) error {
		c.lastPing.Store(time.Now().Unix())
		go func() {
			err := c.SendRawMsg(context.Background(), websocket.PongMessage, nil)
			if err != nil {
				zap.L().Warn("failed to send pong message", zap.Error(err))
			}
		}()
		return nil
	})
	conn.SetCloseHandler(func(code int, text string) error {
		zap.L().Info("websocket connection closed", zap.Int("code", code), zap.String("text", text))
		c.Close()
		return nil
	})
	go c.handlerWriteMessage()
	return c
}
